package com.wave.wavesystem.ai.test.docquery;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.rag.Query;
import org.springframework.ai.rag.preretrieval.query.transformation.QueryTransformer;
import org.springframework.ai.rag.preretrieval.query.transformation.RewriteQueryTransformer;
import org.springframework.stereotype.Component;

/**
 * 查询重新器
 */
@Component
public class MyQueryRewriter {

    private final QueryTransformer queryTransformer;

    public MyQueryRewriter(ChatModel chatModel) {

        queryTransformer = RewriteQueryTransformer.builder()
                .chatClientBuilder(ChatClient.builder(chatModel))
                .build();
    }

    public String doQuery(String query) {
        Query query1 = new Query(query);

        Query transform = queryTransformer.transform(query1);
        return transform.text();
    }


}
